from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.datasets import fetch_20newsgroups
from sklearn.metrics import accuracy_score

# https://blog.csdn.net/weixin_47570444/article/details/141247369
def nbcls():
    # 加载数据集
    news = fetch_20newsgroups(subset='all')
    # 切分数据集
    x_train, x_test, y_train, y_test = train_test_split(news.data,news.target,test_size=0.3,random_state=33)
    # 特征工程
    tf = TfidfVectorizer()
    x_train =  tf.fit_transform(x_train)
    x_test =  tf.transform(x_test)
    # 创建模型
    mlb = MultinomialNB()
    # 训练模型
    mlb.fit(x_train,y_train)

    # 测试
    y_pred = mlb.predict(x_test)


    # 模型评估
    print("预测准确率为：", mlb.score(x_test, y_test))
    print(accuracy_score(y_test, y_pred))

if __name__ == '__main__':
    nbcls()